﻿#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <linux/ip.h>
#include <linux/tcp.h>


#include "packet.h"
#include "log.h"

#define IP4_HDRLEN 20
#define TCP_HDRLEN 20

uint32_t djb2_hash(const void* data, size_t len) {
    const uint8_t* p = (const uint8_t*)data;
    uint32_t hash = 5381;
    for (size_t i = 0; i < len; i++) {
        hash = ((hash << 5) + hash) + p[i]; /* hash * 33 + c */
    }
    return hash;
}


// 计算校验和
unsigned short checksum(unsigned short* buf, size_t  nwords) {
    unsigned long sum;
    for (sum = 0; nwords > 0; nwords--)
        sum += *buf++;
    sum = (sum >> 16) + (sum & 0xffff);
    sum += (sum >> 16);
    return  (unsigned short)~sum;
}

// 计算TCP校验和
unsigned short tcp_checksum(struct iphdr* ip, struct tcphdr* tcp, size_t tcp_len) {
    struct pseudo_header {
        u_int32_t src_address;
        u_int32_t dst_address;
        u_int8_t zero;
        u_int8_t proto;
        u_int16_t length;
    } pseudo_header;

    pseudo_header.src_address = ip->saddr;
    pseudo_header.dst_address = ip->daddr;
    pseudo_header.zero = 0;
    pseudo_header.proto = IPPROTO_TCP;
    pseudo_header.length = htons((u_int16_t)tcp_len);

    // 计算该函数需要的总字节数  
    size_t total_len = tcp_len + sizeof(pseudo_header);
    char* buf = (char*)malloc(total_len + 1); // +1用于填充字节  
    if (buf == NULL) {
        perror("malloc failed");
        return 0; // 内存分配失败处理  
    }

    // 复制伪头部和 TCP 数据  
    memcpy(buf, &pseudo_header, sizeof(pseudo_header));
    memcpy(buf + sizeof(pseudo_header), tcp, tcp_len);

    // 如果总长度是奇数，则添加一个填充字节  
    if (total_len % 2 != 0) {
        buf[total_len] = 0; // 填充字节；  
        total_len++; // 更新总长度  
    }

    // 计算校验和  
    unsigned short result = checksum((unsigned short*)buf, total_len / 2);
    free(buf);
    return result;
}
void tcp_checksum_buff(char * data, size_t total_len) {

    struct iphdr* ip_header = (struct iphdr*)data;;
    

    struct tcphdr* tcp_header = (struct tcphdr*)(data + IP4_HDRLEN);

    size_t data_len = total_len - IP4_HDRLEN - TCP_HDRLEN;


    // 更新 IP 头部的总长度
    ip_header->tot_len = htons((uint16_t)(IP4_HDRLEN + TCP_HDRLEN + data_len));
    ip_header->check = 0;
    ip_header->check = checksum((unsigned short*)ip_header, TCP_HDRLEN / 2);


    // 计算 TCP 校验和，包括数据部分
    tcp_header->check = 0;
    tcp_header->check = tcp_checksum(ip_header, tcp_header, TCP_HDRLEN + data_len);

}


// 构建IP头
struct iphdr* create_ip_header(char* source_ip, char* dest_ip) {
    struct iphdr* ip_header = (struct iphdr*)malloc(IP4_HDRLEN);
    memset(ip_header, 0, IP4_HDRLEN);

    ip_header->version = 4;
    ip_header->ihl = 5;
    ip_header->tos = 0;
    ip_header->tot_len = htons(IP4_HDRLEN + TCP_HDRLEN); // IP头长度 + TCP头长度
    ip_header->id = htons(54321);
    ip_header->frag_off = 0;
    ip_header->ttl = 64;
    ip_header->protocol = IPPROTO_TCP;
    ip_header->check = 0; // 校验和将由内核计算
    ip_header->saddr = inet_addr(source_ip);
    ip_header->daddr = inet_addr(dest_ip);

    // 计算IP校验和
    ip_header->check = checksum((unsigned short*)ip_header, IP4_HDRLEN / 2);

    return ip_header;
}

// 构建TCP头
struct tcphdr* create_tcp_header(uint16_t source_port, uint16_t dest_port, uint32_t seq, uint32_t ack, int flags) {
    struct tcphdr* tcp_header = (struct tcphdr*)malloc(TCP_HDRLEN);
    memset(tcp_header, 0, TCP_HDRLEN);

    tcp_header->source = htons(source_port);
    tcp_header->dest = htons(dest_port);
    tcp_header->seq = htonl(seq);
    tcp_header->ack_seq = htonl(ack);
    tcp_header->doff = 5; // 数据偏移，即TCP头长度
    tcp_header->fin = (flags & 1) ? 1 : 0;
    tcp_header->syn = (flags & 2) ? 1 : 0;
    tcp_header->rst = (flags & 4) ? 1 : 0;
    tcp_header->psh = (flags & 8) ? 1 : 0;
    tcp_header->ack = (flags & 16) ? 1 : 0;
    tcp_header->urg = (flags & 32) ? 1 : 0;
    tcp_header->window = htons(5840);
    tcp_header->check = 0; // 校验和将由内核计算
    tcp_header->urg_ptr = 0;

    return tcp_header;
}


void packet_syn_ack(int tun_fd, SocketNode* socket_node) {

    struct iphdr* ip_header = create_ip_header(socket_node->src_ip, socket_node->dst_ip);
    struct tcphdr* tcp_header = create_tcp_header(socket_node->src_port, socket_node->dst_port, socket_node->seq, socket_node->ack, 2 | 16); 

    // 计算TCP校验和
    tcp_header->check = tcp_checksum(ip_header, tcp_header, sizeof(struct tcphdr));

    char buffer[IP4_HDRLEN + TCP_HDRLEN];
    memcpy(buffer, ip_header, IP4_HDRLEN);
    memcpy(buffer + IP4_HDRLEN, tcp_header, TCP_HDRLEN);


    ssize_t bytes_sent = write(tun_fd, buffer, IP4_HDRLEN + TCP_HDRLEN);
    if (bytes_sent < 0) {
        perror("Sendto failed");
        close(tun_fd);
    }
    else if (bytes_sent != IP4_HDRLEN + TCP_HDRLEN) {
        fprintf(stderr, "Incomplete send: sent %zd bytes\n", bytes_sent);
        close(tun_fd);
        exit(EXIT_FAILURE);
    }


    free(ip_header);
    free(tcp_header);

}

void packet_ack(int tun_fd, SocketNode* socket_node) {

    struct iphdr* ip_header = create_ip_header(socket_node->src_ip, socket_node->dst_ip);
    struct tcphdr* tcp_header = create_tcp_header(socket_node->src_port, socket_node->dst_port, socket_node->seq, socket_node->ack, 16);
    // 计算TCP校验和
    tcp_header->check = tcp_checksum(ip_header, tcp_header, sizeof(struct tcphdr));

    char buffer[IP4_HDRLEN + TCP_HDRLEN];
    memcpy(buffer, ip_header, IP4_HDRLEN);
    memcpy(buffer + IP4_HDRLEN, tcp_header, TCP_HDRLEN);


    ssize_t bytes_sent = write(tun_fd, buffer, IP4_HDRLEN + TCP_HDRLEN);
    if (bytes_sent < 0) {
        perror("Sendto failed");
        close(tun_fd);
    }
    else if (bytes_sent != IP4_HDRLEN + TCP_HDRLEN) {
        fprintf(stderr, "Incomplete send: sent %zd bytes\n", bytes_sent);
        close(tun_fd);
        exit(EXIT_FAILURE);
    }


    free(ip_header);
    free(tcp_header);

}


void packet_push(int tun_fd, SocketNode* socket_node, char* data, size_t data_len) {
    struct iphdr* ip_header = create_ip_header(socket_node->src_ip, socket_node->dst_ip);
    if (ip_header == NULL) {
        perror("create_ip_header failed");
        close(tun_fd);
        exit(EXIT_FAILURE);
    }

    struct tcphdr* tcp_header = create_tcp_header(socket_node->src_port, socket_node->dst_port, socket_node->seq, socket_node->ack, 8 | 16); // PSH | ACK
    if (tcp_header == NULL) {
        perror("create_tcp_header failed");
        free(ip_header);
        close(tun_fd);
        exit(EXIT_FAILURE);
    }

    // 更新 IP 头部的总长度
    ip_header->tot_len = htons((uint16_t)(IP4_HDRLEN + TCP_HDRLEN + data_len));
    ip_header->check = 0;
    ip_header->check = checksum((unsigned short*)ip_header, IP4_HDRLEN / 2);


    // 创建缓冲区并复制数据
    char buffer[IP4_HDRLEN + TCP_HDRLEN + data_len];
    memcpy(buffer, ip_header, sizeof(struct iphdr));
    memcpy(buffer + IP4_HDRLEN, tcp_header, TCP_HDRLEN);
    memcpy(buffer + IP4_HDRLEN + TCP_HDRLEN, data, data_len);


    tcp_checksum_buff(buffer, IP4_HDRLEN + TCP_HDRLEN + data_len);

    log_debug_print(buffer, IP4_HDRLEN + TCP_HDRLEN + data_len);

    // 发送数据包
    ssize_t bytes_sent = write(tun_fd, buffer, IP4_HDRLEN + TCP_HDRLEN + data_len);
    if (bytes_sent < 0) {
        perror("Sendto failed");
        free(ip_header);
        free(tcp_header);
        close(tun_fd);
        exit(EXIT_FAILURE);
    }
    else if (bytes_sent != IP4_HDRLEN + TCP_HDRLEN + data_len) {
        fprintf(stderr, "Incomplete send: sent %zd bytes\n", bytes_sent);
        free(ip_header);
        free(tcp_header);
        close(tun_fd);
        exit(EXIT_FAILURE);
    }

    free(ip_header);
    free(tcp_header);
}

void packet_fin(int tun_fd, SocketNode* socket_node) {
    struct iphdr* ip_header = create_ip_header(socket_node->src_ip, socket_node->dst_ip);
    struct tcphdr* tcp_header = create_tcp_header(socket_node->src_port, socket_node->dst_port, socket_node->seq, socket_node->ack, 1 | 16); // fin | ACK
    // 计算TCP校验和
    tcp_header->check = tcp_checksum(ip_header, tcp_header, sizeof(struct tcphdr));

    char buffer[IP4_HDRLEN + TCP_HDRLEN];
    memcpy(buffer, ip_header, IP4_HDRLEN);
    memcpy(buffer + IP4_HDRLEN, tcp_header, TCP_HDRLEN);


    ssize_t bytes_sent = write(tun_fd, buffer, IP4_HDRLEN + TCP_HDRLEN);
    if (bytes_sent < 0) {
        perror("Sendto failed");
        close(tun_fd);
    }
    else if (bytes_sent != IP4_HDRLEN + TCP_HDRLEN) {
        fprintf(stderr, "Incomplete send: sent %zd bytes\n", bytes_sent);
        close(tun_fd);
        exit(EXIT_FAILURE);
    }


    free(ip_header);
    free(tcp_header);
}

void packet_rst(int tun_fd, SocketNode* socket_node) {

    struct iphdr* ip_header = create_ip_header(socket_node->src_ip, socket_node->dst_ip);
    struct tcphdr* tcp_header = create_tcp_header(socket_node->src_port, socket_node->dst_port, socket_node->seq, socket_node->ack, 4);
    // 计算TCP校验和
    tcp_header->check = tcp_checksum(ip_header, tcp_header, sizeof(struct tcphdr));

    char buffer[IP4_HDRLEN + TCP_HDRLEN];
    memcpy(buffer, ip_header, IP4_HDRLEN);
    memcpy(buffer + IP4_HDRLEN, tcp_header, TCP_HDRLEN);


    ssize_t bytes_sent = write(tun_fd, buffer, IP4_HDRLEN + TCP_HDRLEN);
    if (bytes_sent < 0) {
        perror("Sendto failed");
        close(tun_fd);
    }
    else if (bytes_sent != IP4_HDRLEN + TCP_HDRLEN) {
        fprintf(stderr, "Incomplete send: sent %zd bytes\n", bytes_sent);
        close(tun_fd);
        exit(EXIT_FAILURE);
    }


    free(ip_header);
    free(tcp_header);

}


